package com.zheng.rpc.remoting.transport.netty.client;

import com.zheng.rpc.enums.CompressTypeEnum;
import com.zheng.rpc.enums.SerializationTypeEnum;
import com.zheng.rpc.extensions.ExtensionLoader;
import com.zheng.rpc.factories.SingletonFactory;
import com.zheng.rpc.registry.ServiceDiscovery;
import com.zheng.rpc.remoting.constants.RpcConstants;
import com.zheng.rpc.remoting.dto.RpcMessage;
import com.zheng.rpc.remoting.dto.RpcRequest;
import com.zheng.rpc.remoting.dto.RpcResponse;
import com.zheng.rpc.remoting.transport.RpcRequestTransport;
import com.zheng.rpc.remoting.transport.netty.codec.RpcMessageDecoder;
import com.zheng.rpc.remoting.transport.netty.codec.RpcMessageEncoder;
import io.netty.bootstrap.Bootstrap;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.logging.LogLevel;
import io.netty.handler.logging.LoggingHandler;
import io.netty.handler.timeout.IdleStateHandler;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

@Slf4j
public class NettyRpcClient implements RpcRequestTransport {
    private final ServiceDiscovery serviceDiscovery;
    private final UnprocessedRequests unprocessedRequests;
    private final ChannelProvider channelProvider;
    private final Bootstrap bootstrap;
    private final EventLoopGroup eventLoopGroup;

    public NettyRpcClient() {
        eventLoopGroup = new NioEventLoopGroup();
        bootstrap = new Bootstrap();
        bootstrap.group(eventLoopGroup)
                .channel(NioSocketChannel.class)
                .handler(new LoggingHandler(LogLevel.INFO))
                .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 5000)
                .handler(new ChannelInitializer<SocketChannel>() {
                    @Override
                    protected void initChannel(SocketChannel ch) throws Exception {
                        ChannelPipeline pipeline = ch.pipeline();
                        // 如果15秒钟内没有数据发送到服务器，则发送心跳请求
                        pipeline.addLast(new IdleStateHandler(0, 5, 0, TimeUnit.SECONDS));
                        pipeline.addLast(new RpcMessageEncoder());
                        pipeline.addLast(new RpcMessageDecoder());
                        pipeline.addLast(new NettyRpcClientHandler());
                    }
                });
        this.serviceDiscovery = ExtensionLoader.getExtensionLoader(ServiceDiscovery.class).getExtension("zk");
        this.unprocessedRequests = SingletonFactory.getInstance(UnprocessedRequests.class);
        this.channelProvider = SingletonFactory.getInstance(ChannelProvider.class);
    }

    /**
     * 连接服务器并获取channel，就可以向服务器发送rpc消息
     *
     * @param address server address
     * @return the channel
     */
    @SneakyThrows
    public Channel doConnect(InetSocketAddress address) {
        CompletableFuture<Channel> completableFuture = new CompletableFuture<>();
        bootstrap.connect(address).addListener((ChannelFutureListener)future -> {
            if (future.isSuccess()) {
                log.info("The client has connected [{}] successful!", address.toString());
                completableFuture.complete(future.channel());
            } else {
                throw new IllegalStateException();
            }
        });
        return completableFuture.get();
    }

    /**
     * 传输 rpc 请求(RpcRequest) 到服务端。
     * @param rpcRequest message body
     * @return
     */
    @Override
    public Object sendRpcRequest(RpcRequest rpcRequest) {
        CompletableFuture<RpcResponse<Object>> resultFuture = new CompletableFuture<>();
        String serviceName = rpcRequest.toRpcProperties().toRpcServiceName();
        InetSocketAddress address = serviceDiscovery.lookupService(serviceName);
        Channel channel = getChannel(address);
        if (channel.isActive()) {
            // 添加未处理请求
            unprocessedRequests.put(rpcRequest.getRequestId(), resultFuture);
            RpcMessage message = new RpcMessage();
            message.setData(rpcRequest);
            message.setCodec(SerializationTypeEnum.PROTOSTUFF.getCode());
            message.setCompress(CompressTypeEnum.GZIP.getCode());
            message.setMessageType(RpcConstants.REQUEST_TYPE);
            channel.writeAndFlush(message).addListener((ChannelFutureListener)future -> {
                if (future.isSuccess()) {
                    log.info("client send message: [{}]", message);
                } else {
                    future.channel().close();
                    resultFuture.completeExceptionally(future.cause());
                    log.error("Send failed" + future.cause());
                }
            });
        } else {
            throw new IllegalStateException();
        }
        return resultFuture;
    }

    /**
     * 连接服务端（目标方法所在的服务器）并返回对应的 Channel。
     * @param address
     * @return
     */
    public Channel getChannel(InetSocketAddress address) {
        Channel channel = channelProvider.get(address);
        if (channel == null) {
            channel = doConnect(address);
            channelProvider.set(address, channel);
        }
        return channel;
    }

    public void close() {
        eventLoopGroup.shutdownGracefully();
    }
}
